/*
 * @BEGIN LICENSE
 *
 * Psi4: an open-source quantum chemistry software package
 *
 * Copyright (c) 2007-2024 The Psi4 Developers.
 *
 * The copyrights for code used from other parties are included in
 * the corresponding files.
 *
 * This file is part of Psi4.
 *
 * Psi4 is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * Psi4 is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License along
 * with Psi4; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * @END LICENSE
 */

/*!
  \file
  \brief Obtain orbital space and reordering for CI/MCSCF wavefunctions
  \ingroup QT
*/

#include "qt.h"

#include "psi4/libciomr/libciomr.h"
#include "psi4/psifiles.h"
#include "psi4/liboptions/liboptions.h"
#include "psi4/libpsi4util/PsiOutStream.h"

#include <cstdio>
#include <cstdlib>

namespace psi {

/*!
** ras_set3()
**
** Updated version of ras_set2() to conform with new DETCI naming of
** various orbital spaces, particularly to distinguish between restricted
** orbitals (those that are doubly occupied but can optimized in an MCSCF)
** and frozen orbitals (those that are doubly occupied but cannot be
** optimized in an MCSCF).  Also remove some deprecated code supporting
** Mark Hoffmann ordering, etc.
**
** This function sets up the number of orbitals per irrep for each of the
** RAS subspaces [frozen core, restricted core, RAS I, RAS II,
** RAS III, RAS IV, restricted virts, frozen virts].
** It also obtains the appropriate orbital reordering array.  The
** reordering array takes a basis function in Pitzer ordering (orbitals
** grouped according to irrep) and gives the corresponding index
** in the RAS numbering scheme.  Orbitals are numbered according to
** irrep within each of the subspaces.
**
** Guesses for docc and socc should be provided, and they will be left
** as-is if they are not present in input (maybe auto-guessed by program)
**
** C. David Sherrill
**
**  \param nirreps     =  num of irreps in computational point group
**  \param nmo         =  number of MO's
**  \param orbspi      =  array giving num symmetry orbitals (or MOs) per irrep
**  \param docc        =  array of doubly occupied orbitals per irrep
**                        (guess must be provided)
**  \param socc        =  array of singly occupied orbitals per irrep
**                        (guess must be provided)
**  \param frdocc      =  array of frozen core per irrep
**                        (returned by function, but allocate before call)
**  \param fruocc      =  array of frozen virtuals per irrep
**                        (returned by function, but allocate before call)
**  \param rstrdocc    =  array of restricted core per irrep
**                        (returned by function, but allocate before call)
**  \param rstruocc    =  array of restricted unoccupied per irrep
**                        (returned by function, but allocate before call)
**  \param ras_opi     =  matrix giving num of orbitals per irrep per ras space,
**                        addressed as ras_opi[ras_space][irrep]
**                        (returned by function, but allocate before call)
**                        RAS IV will not be deduced unless RAS III is
**                        present in user input (in which case RAS IV needs
**                        to make up any remaining orbitals)
**  \param core_guess  =  array of the core orbitals per irrep to drop
**                        (must be provided).  This is copied into
**                        FROZEN_DOCC (if is_mcscf == false) or
**                        RESTRICTED_DOCC (if is_mcscf == true), if both
**                        of those two arrays are absent from input
**  \param order       =  array nmo big which maps Pitzer to Correlated order
**                        (returned by function, but allocate before call)
**  \param ras_type    =  if 1, put docc and socc together in same RAS space
**                        (RAS I), as appropriate for DETCI.  If 0, put socc
**                        in its own RAS space (RAS II), as appropriate for CC.
**  \param is_mcscf    =  true if an MCSCF-type wavefunction (orbital
**                        optimization), else false.  Used to determine
**                        default behavior of FROZEN_DOCC, RESTRICTED_DOCC,
**                        FROZEN_UOCC, RESTRICTED_UOCC
**  \param options     =  Options object used to parse user input
**
** Rules for how spaces can be specified:
**   1. DOCC and SOCC are not changed by this routine
**      (guesses might have been auto-generated by Psi) unless explicitly
**      given in user input, in which case the guess arrays are overwritten
**      by the user input in this routine
**   2. FROZEN_DOCC and RESTRICTED_DOCC will be as given by the user.  If
**      the user does not specify either of these keywords, then maybe
**      they specified FREEZE_CORE = TRUE instead, and we want to support
**      that.  If neither FROZEN_DOCC nor RESTRICTED_DOCC is specified, then
**      we will use the array core_guess[] to get them.  This array is
**      obtained from whatever routine figures out how many orbitals per
**      irrep to freeze when FREEZE_CORE = true (passed down as an argument
**      into this function).  If the computation is an MCSCF, then we'll
**      copy core_guess into RESTRICTED_DOCC.  If it's not an MCSCF, then
**      we'll copy it into FROZEN_DOCC.
**   3. The user can give either ACTIVE (for a CAS type computation, where
**      internally this is basically like RAS 2), or RAS keywords.  Not both.
**   4. By default, unused virtual orbitals go into FROZEN_UOCC unless
**      we're told this is an MCSCF, in which case they should go by
**      default into RESTRICTED_UOCC
**   5. RAS 4 is never used unless explicitly given in input; it's a special
**      feature not normally desired.  (Note: this is a change from previous
**      default behavior in which unused orbitals were stuffed into RAS 4)
**
** Returns: 1 for success, 0 otherwise
** \ingroup QT
*/
int ras_set3(int nirreps, int nmo, int *orbspi, int *docc, int *socc, int *frdocc, int *fruocc, int *restrdocc,
             int *restruocc, int **ras_opi, int *core_guess, int *order, int ras_type, bool is_mcscf,
             Options &options) {
    int i, irrep, point, tmpi, cnt = 0;
    int errcod, errbad = 0;
    int *used, *offset, **tras;
    int *tmp_frdocc, *tmp_fruocc;
    bool parsed_ras1 = false, parsed_ras2 = false;
    bool parsed_ras3 = false, parsed_ras4 = false;
    bool parsed_frozen_docc = false, parsed_restr_docc = false;
    bool parsed_frozen_uocc = false, parsed_restr_uocc = false;

    used = init_int_array(nirreps);
    offset = init_int_array(nirreps);

    // zero the things where guesses shouldn't be coming in
    for (i = 0; i < MAX_RAS_SPACES; i++) {
        zero_int_array(ras_opi[i], nirreps);
    }
    zero_int_array(frdocc, nirreps);
    zero_int_array(restrdocc, nirreps);
    zero_int_array(fruocc, nirreps);
    zero_int_array(restruocc, nirreps);

    zero_int_array(order, nmo);

    // Replace DOCC and SOCC only if they are in input, otherwise just
    // leave alone the guess provided.
    if (options["DOCC"].has_changed()) {
        if (options["DOCC"].size() != nirreps) {
            throw InputException("ras_set2(): Wrong size of array", "DOCC", __FILE__, __LINE__);
        }
        options.fill_int_array("DOCC", docc);
    }
    if (options["SOCC"].has_changed()) {
        if (options["SOCC"].size() != nirreps) {
            throw InputException("ras_set2(): Wrong size of array", "SOCC", __FILE__, __LINE__);
        }
        options.fill_int_array("SOCC", socc);
    }

    // Take FROZEN_DOCC from user input if it is there
    if (options["FROZEN_DOCC"].has_changed()) {
        if (options["FROZEN_DOCC"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "FROZEN_DOCC", __FILE__, __LINE__);
        }
        options.fill_int_array("FROZEN_DOCC", frdocc);
        parsed_frozen_docc = true;
    }

    // Take RESTRICTED_DOCC from user input if it is there
    if (options["RESTRICTED_DOCC"].has_changed()) {
        if (options["RESTRICTED_DOCC"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "RESTRICTED_DOCC", __FILE__, __LINE__);
        }
        options.fill_int_array("RESTRICTED_DOCC", restrdocc);
        parsed_restr_docc = true;
    }

    // If we weren't given either FROZEN_DOCC *or* RESTRICTED_DOCC, then
    // we'd better guess something.  Take default from core_guess[], and
    // put that in either in frdocc (if not an MCSCF) or restrdocc (if an MCSCF)
    if (!parsed_frozen_docc && !parsed_restr_docc) {
        if (!is_mcscf) {
            for (i = 0; i < nirreps; i++) {
                frdocc[i] = core_guess[i];
            }
        } else {
            for (i = 0; i < nirreps; i++) {
                restrdocc[i] = core_guess[i];
            }
        }
    }
    // at this point, we have FROZEN_DOCC and RESTRICTED_DOCC

    // Change FROZEN_UOCC from zero array if user has it in input.
    // If no user input detected, we might deduce a
    // different value based on the values of other spaces (below)
    if (options["FROZEN_UOCC"].has_changed()) {
        if (options["FROZEN_UOCC"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "FROZEN_UOCC", __FILE__, __LINE__);
        }
        options.fill_int_array("FROZEN_UOCC", fruocc);
        parsed_frozen_uocc = true;
    }

    // Change RESTRICTED_UOCC from zero array if user has it in input.
    // If no user input detected, we might deduce a
    // different value based on the values of other spaces (below)
    if (options["RESTRICTED_UOCC"].has_changed()) {
        if (options["RESTRICTED_UOCC"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "RESTRICTED_UOCC", __FILE__, __LINE__);
        }
        options.fill_int_array("RESTRICTED_UOCC", restruocc);
        parsed_restr_uocc = true;
    }

    // Now use the parser to get the arrays we require IF they are
    // present in user input ... otherwise leave them alone (they just got
    // filled with zeroes above, so that will be the default until we
    // try do deduce them from other information below).
    // Set some flags for some of the spaces we have parsed to signal we
    // don't need to deduce those spaces.

    if (options["RAS1"].has_changed()) {
        if (options["RAS1"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "RAS1", __FILE__, __LINE__);
        }
        options.fill_int_array("RAS1", ras_opi[0]);
        parsed_ras1 = true;
    }
    if (options["RAS2"].has_changed()) {
        if (options["RAS2"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "RAS2", __FILE__, __LINE__);
        }
        options.fill_int_array("RAS2", ras_opi[1]);
        parsed_ras2 = true;
    }
    if (options["RAS3"].has_changed()) {
        if (options["RAS3"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "RAS3", __FILE__, __LINE__);
        }
        options.fill_int_array("RAS3", ras_opi[2]);
        parsed_ras3 = true;
    }
    if (options["RAS4"].has_changed()) {
        if (options["RAS4"].size() != nirreps) {
            throw InputException("ras_set3(): Wrong size of array", "RAS4", __FILE__, __LINE__);
        }
        options.fill_int_array("RAS4", ras_opi[3]);
        parsed_ras4 = true;
    }

    // If the user has not specified RAS 1, we must deduce it.
    // As of 2004, RAS 1 will not include any frozen core or restricted core
    // orbitals.

    if (!parsed_ras1) {
        for (irrep = 0; irrep < nirreps; irrep++) {
            if (ras_type == 1) ras_opi[0][irrep] = docc[irrep] + socc[irrep];
            if (ras_type == 0) ras_opi[0][irrep] = docc[irrep];
            ras_opi[0][irrep] -= (frdocc[irrep] + restrdocc[irrep]);
        }
    }
    // at this point, we now have an array in RAS 1

    // If the user hasn't specified RAS-style keywords, look for ACTIVE.
    // Assume this always goes after any frozen core or restricted core.

    if (parsed_ras1 || parsed_ras2 || parsed_ras3 || parsed_ras4) {
        if (options["ACTIVE"].has_changed()) {
            throw InputException("ras_set3(): RAS and ACTIVE keywords mutually exclusive", "ACTIVE", __FILE__,
                                 __LINE__);
        }
    }

    // Next we need to provide a guess for RAS 2
    if (!parsed_ras2) {  // we didn't find a RAS keyword so look for "ACTIVE"
        if (options["ACTIVE"].has_changed()) {
            if (options["ACTIVE"].size() != nirreps) {
                throw InputException("ras_set3(): Wrong size of array", "ACTIVE", __FILE__, __LINE__);
            }
            options.fill_int_array("ACTIVE", ras_opi[1]);  // fill RAS 2 with ACTIVE

            // ACTIVE overrides RAS 1 ... any occupieds not in RAS 2 should
            // be restricted or frozen core, not RAS 1
            for (irrep = 0; irrep < nirreps; irrep++) ras_opi[0][irrep] = 0;

            // if not MCSCF, default remaining virts to FROZEN_UOCC
            if (!is_mcscf && !parsed_frozen_uocc) {
                for (irrep = 0; irrep < nirreps; irrep++) {
                    fruocc[irrep] = orbspi[irrep] - frdocc[irrep] - restrdocc[irrep] - ras_opi[0][irrep] -
                                    ras_opi[1][irrep] - restruocc[irrep];
                }
                parsed_frozen_uocc = true;
            }
            // if it is MCSCF, default remaining virts to RESTRICTED_UOCC
            if (is_mcscf && !parsed_restr_uocc) {
                for (irrep = 0; irrep < nirreps; irrep++) {
                    restruocc[irrep] = orbspi[irrep] - frdocc[irrep] - restrdocc[irrep] - ras_opi[0][irrep] -
                                       ras_opi[1][irrep] - fruocc[irrep];
                }
                parsed_restr_uocc = true;
            }

            // do a quick check
            for (irrep = 0; irrep < nirreps; irrep++) {
                if (frdocc[irrep] + restrdocc[irrep] + ras_opi[0][irrep] + ras_opi[1][irrep] <
                    docc[irrep] + socc[irrep])
                    outfile->Printf("ras_set3():Warning:Occupied electrons beyond ACTIVE orbs!\n");
            }
        }       // end case where we found ACTIVE keyword
        else {  // we didn't find ACTIVE keyword, so we need to determine RAS 2
            for (irrep = 0; irrep < nirreps; irrep++) {
                if (ras_type == 1) ras_opi[1][irrep] = 0;
                if (ras_type == 0) ras_opi[1][irrep] = socc[irrep];
            }
        }
    }  // at this point, we now have an array in RAS 2

    // Now we need to figure out RAS 3
    // New for 2015: RAS 4 will not be used unless explicitly specified.
    if (!parsed_ras3) {
        for (irrep = 0; irrep < nirreps; irrep++) {
            tmpi = frdocc[irrep] + restrdocc[irrep] + fruocc[irrep] + restruocc[irrep] + ras_opi[0][irrep] +
                   ras_opi[1][irrep] + ras_opi[3][irrep];
            if (tmpi > orbspi[irrep]) {
                outfile->Printf("ras_set3(): orbitals don't add up for irrep %d\n", irrep);
                return (0);
            }
            ras_opi[2][irrep] = orbspi[irrep] - tmpi;
        }
    }  // at this point, we now have an array in RAS 3

    // Now if FROZEN_UOCC or RESTRICTED_UOCC have not been set yet, fill
    // the appropriate one with any unused orbitals (which one depends on
    // whether or not this is an MCSCF)
    if (!is_mcscf) {  // if not an MCSCF, default into fruocc
        if (!parsed_frozen_uocc) {
            for (irrep = 0; irrep < nirreps; irrep++) {
                tmpi = frdocc[irrep] + restrdocc[irrep] + fruocc[irrep] + restruocc[irrep] + ras_opi[0][irrep] +
                       ras_opi[1][irrep] + ras_opi[2][irrep] + ras_opi[3][irrep];
                if (tmpi > orbspi[irrep]) {
                    outfile->Printf("ras_set3(): orbitals don't add up for irrep %d\n", irrep);
                    return (0);
                }
                fruocc[irrep] = orbspi[irrep] - tmpi;
            }
        }
    } else {  // is an MCSCF, so default into restruocc instead of fruocc
        if (!parsed_restr_uocc) {
            for (irrep = 0; irrep < nirreps; irrep++) {
                tmpi = frdocc[irrep] + restrdocc[irrep] + fruocc[irrep] + restruocc[irrep] + ras_opi[0][irrep] +
                       ras_opi[1][irrep] + ras_opi[2][irrep] + ras_opi[3][irrep];
                if (tmpi > orbspi[irrep]) {
                    outfile->Printf("ras_set3(): orbitals don't add up for irrep %d\n", irrep);
                    return (0);
                }
                restruocc[irrep] = orbspi[irrep] - tmpi;
            }
        }
    }  // should complete fruocc and/or restruocc

    // copy everything to some temporary arrays just to count easier
    tras = init_int_matrix(MAX_RAS_SPACES + 4, nirreps);

    for (irrep = 0; irrep < nirreps; irrep++) {
        tras[0][irrep] = frdocc[irrep];
        tras[1][irrep] = restrdocc[irrep];
        tras[MAX_RAS_SPACES + 2][irrep] = restruocc[irrep];
        tras[MAX_RAS_SPACES + 3][irrep] = fruocc[irrep];
    }
    for (i = 0; i < MAX_RAS_SPACES; i++) {
        for (irrep = 0; irrep < nirreps; irrep++) {
            tras[i + 2][irrep] = ras_opi[i][irrep];
        }
    }

    // construct the offset array
    offset[0] = 0;
    for (irrep = 1; irrep < nirreps; irrep++) {
        offset[irrep] = offset[irrep - 1] + orbspi[irrep - 1];
    }

    for (i = 0; i < MAX_RAS_SPACES + 4; i++) {
        for (irrep = 0; irrep < nirreps; irrep++) {
            while (tras[i][irrep]) {
                point = used[irrep] + offset[irrep];
                // outfile->Printf("%d %d %d %d %d %d %d \n", i, irrep, tras[i][irrep], used[irrep], offset[irrep],
                // point, nmo);
                if (point < 0 || point >= nmo) {
                    throw PsiException("ras_set2(): Invalid point value", __FILE__, __LINE__);
                }
                order[point] = cnt++;
                used[irrep]++;
                tras[i][irrep]--;
            }
        }
    }

    // do a final check
    for (irrep = 0; irrep < nirreps; irrep++) {
        if (used[irrep] > orbspi[irrep]) {
            outfile->Printf("ras_set3(): on final check, used more orbitals");
            outfile->Printf("   than were available (%d vs %d) for irrep %d\n", used[irrep], orbspi[irrep], irrep);
            errbad = 1;
        }
        if (used[irrep] < orbspi[irrep]) {
            outfile->Printf("ras_set3(): on final check, used fewer orbitals");
            outfile->Printf("   than were available (%d vs %d) for irrep %d\n", used[irrep], orbspi[irrep], irrep);
            errbad = 1;
        }
    }

    free(used);
    free(offset);
    free_int_matrix(tras);

    return (!errbad);
}
}  // namespace psi
